package com.beney.ds.trees.avl;

public class AVLTree<K extends Comparable<K>, V> {
    private Node root;

    private class Node {
        K key;
        V val;
        int height;
        Node left, right;

        public Node(K key, V val, Node left, Node right) {
            this.key = key;
            this.val = val;
            this.left = left;
            this.right = right;
        }
    }

    public int height() {
        return height(root);
    }

    private int height(Node x) {
        if (x == null)
            return -1;
        return x.height;
    }

    /**
     * @param key 节点匹配关键码
     * @return 根据key获取的节点值，若节点不存在返回null
     */
    public V get(K key) {
        if (key == null)
            return null;
        Node x = get(root, key);
        return x == null ? null : x.val;
    }

    /**
     * BST搜索
     *
     * @param x 当前遍历到的节点
     * @param key 要查找的key
     * @return 匹配节点。失败则null
     */
    private Node get(Node x, K key) {
        if (x == null)
            return null;
        int cmp = key.compareTo(x.key);
        if (cmp < 0)        return get(x.left, key);
        else if (cmp > 0)   return get(x.right, key);
        else                return x;
    }

    /**
     * @return true 插入成功，false 插入失败
     */
    public boolean put(K k, V v) {
        if (k == null)
            return false;
        root = put(root, k, v);
        return true;
    }

    // 返回在x子树上插入kv之后的树
    private Node put(Node x, K k, V v) {
        if (x == null)
            return new Node(k, v, null, null);
        int cmp = k.compareTo(x.key);
        if (cmp < 0)        x.left = put(x.left, k, v);
        else if (cmp > 0)   x.right = put(x.right, k, v);
        else {
            x.val = v;
            return x;
        }
        x.height = 1 + Math.max(height(x.left), height(x.right));
        return balance(x); // 自底向上进行平衡
    }

    /**
     * @param x 子树的根
     * @return 平衡之后的子树
     */
    private Node balance(Node x) {
        if (x == null) return null;
        int factor = balanceFactor(x);
        if (factor < -1) { // right cases
            if (balanceFactor(x.right) > 0) { // right-left case
                x.right = rotateRight(x.right); // convert to right-right
            }
            x = rotateLeft(x); // right-right case
        }
        else if (factor > 1) { // left cases
            if (balanceFactor(x.left) < 0) { // left-right case
                x.left = rotateLeft(x.left); // convert to left-left
            }
            x = rotateRight(x); // left-left case
        }
        return x;
    }

    /**
     * @param x 子树根
     * @return 左子树高度 - 右子树高度
     */
    private int balanceFactor(Node x) {
        return height(x.left) - height(x.right);
    }

    /**
     * @param x 右旋支点
     * @return 右旋完后子树的根
     */
    private Node rotateRight(Node x) {
        final Node y = x.left;
        x.left = y.right;
        y.right = x;
        x.height = 1 + Math.max(height(x.left), height(x.right));
        y.height = 1 + Math.max(height(y.left), height(y.right));
        return y;
    }

    /**
     * 左旋
     *
     * @param x 左旋支点
     * @return 左旋完后子树的根
     */
    private Node rotateLeft(Node x) {
        final Node y = x.right;
        x.right = y.left;
        y.left = x;
        x.height = 1 + Math.max(height(x.left), height(x.right));
        y.height = 1 + Math.max(height(y.left), height(y.right));
        return y;
    }

    public void delete(K k) {
        if (k == null)
            return;
        root = delete(root, k);
    }

    private Node delete(Node x, K k) {
        int cmp = k.compareTo(x.key);
        if (cmp < 0)        x.left = delete(x.left, k);
        else if (cmp > 0)   x.right = delete(x.right, k);
        else {
            if (x.left == null)
                return x.right;
            else if (x.right == null)
                return x.left;
            else {
                Node oldX = x;
                x = minNode(x.right);
                x.right = deleteMin(oldX.right);
                x.left = oldX.left;
            }
        }
        x.height = 1 + Math.max(height(x.left), height(x.right));
        return balance(x); // 自底向上进行平衡
    }

    /**
     * @param x 子树根
     * @return x为根的子树中的最小节点
     */
    private Node minNode(Node x) {
        Node tmp = null;
        while (x != null) {
            tmp = x;
            x = x.left;
        }
        return tmp;
    }

    /**
     * @param x 子树的根
     * @return 删除最小节点后的树
     */
    private Node deleteMin(Node x) {
        if (x.left == null)
            return x.right;
        x.left = deleteMin(x.left);
        x.height = 1 + Math.max(height(x.left), height(x.right));
        return balance(x);
    }
}
